#!/usr/local/bin/lua
-- process spamfilter commands
-- copyright (C) 2005, 2006 Fidelis Assis <fidelis@pobox.com>

function help()
  local f = io.open(global_osbf_dir .. "spamfilter.help", "r")
  local help_txt = f:read("*all")
  f:close()
  return string.format(help_txt, osbf.cfg_tag_spam, osbf.cfg_tag_unsure_spam,
                       osbf.cfg_tag_unsure_ham, osbf.cfg_tag_ham,
		       osbf.cfg_threshold, osbf.cfg_threshold,
		       osbf.cfg_min_pR_success, osbf._VERSION, sf_version)
end

-- for case insensitive match - from PIL
function nocase (s)
  s = string.gsub(s, "%a", function (c)
        return string.format("[%s%s]", string.lower(c),
                                       string.upper(c))
      end)
  return s
end

-- remove unwanted headers
function clean_header(header, unwanted_headers)
  for _, h in ipairs(unwanted_headers) do
    if (h ~= "") then
      -- h = escape_regchars(h)
      h = nocase(h)
      repeat
        header, count = string.gsub(header, "\n" .. h .. ".-(\n[^ \t])", "%1")
      until count == 0
    end
  end
  return header
end

-- save a list
function list_save(list, listname)
  local file = lists_dir .. listname .. ".lua"
  local w = io.open(file, "w")
  if (w) then
    w:write(listname, " = {")
    local list_contents = ""
    for _, e in ipairs(list) do
      e = string.gsub(e, '\\', '\\\\')
      e = string.gsub(e, '"', '\\"')
      if list_contents == "" then
        list_contents = '"' .. e .. '"'
      else
        list_contents = list_contents .. ',\n"' .. e .. '"'
      end
    end
    if list_contents == '' then
      w:write('}\n')
    else
      w:write('\n', list_contents, '\n}\n')
    end
    w:close()
  else
    return nil, "can't open listname for writing"
  end
  return true
end

-- return a string with the list contents
function list_sprint(list, listname)
  local list_string
  list_string = ""
  if list and listname then
        list_string = "--- start of " .. listname .. " ---\n"
        for _, v in ipairs (list) do
          list_string = list_string .. v .. "\n"
        end
        list_string = list_string .. "--- end of " .. listname .. " ---\n"
  else
  	list_string = nil, "invalid list or listname"
  end
  return  list_string
end

-- insert element into a list (black or white) and save the list
function list_insert(element, list, listname)
  if (element and list and listname) then
    table.insert(list, element)
    -- save the list
    return list_save(list, listname)
  else
    return nil, "invalid argument"
  end
end

-- delete element from a list (black or white) and save the list
function list_delete(element, list, listname)
  local found = false
  if (element and list and listname) then
     for i, e in ipairs(list) do
       if e == element then
          table.remove(list, i)
	  found = true
       end
     end
     -- save the list
     if found then
       return list_save(list, listname)
     else
       return nil, "Element not found in the " .. listname
     end
  else
    return nil, "invalid element or list"
  end
end

-- returns a string report with dbfile statistics
function db_stats (dbfile)
    local version = "OSBF-Bayes"
    local stats_lua = osbf.stats(dbfile)
    if (not stats_lua.version == 5) then
      return "Unknown version"
    end
    
    local report = "-- Statistics for " .. dbfile .. "\n"
    report = report .. string.format(
        "%-35s%12s\n%-35s%12d\n%-35s%12.1f\n%-35s%12d\n%-35s%12d\n%-35s%12d\n",
        "Version:", version,
        "Total buckets in database:", stats_lua.buckets,
        "Buckets used (%):", stats_lua.use * 100,
        "Trainings:", stats_lua.learnings,
        "Bucket size (bytes):", stats_lua.bucket_size,
        "Header size (bytes):", stats_lua.header_size)
    report = report .. string.format("%-35s%12d\n%-35s%12d\n" ..
	"%-35s%12.1f\n%-35s%12d\n%-35s%12d\n",
        "Number of chains:", stats_lua.chains,
        "Max chain len (buckets):", stats_lua.max_chain,
        "Average chain length (buckets):", stats_lua.avg_chain,
        "Max bucket displacement:", stats_lua.max_displacement,
        "Number of unreachable buckets:", stats_lua.unreachable)


    return report
end

-- train "text" as belonging to class "class_index"
function osbf_train(text, class_index)
  local pR, msg_error = osbf.classify(text, osbf.cfg_dbset,
					osbf.cfg_classify_flags)
  if (pR) then
    if ( ( (pR < 0)  and (class_index == osbf.cfg_nonspam_index) ) or
         ( (pR >= 0) and (class_index == osbf.cfg_spam_index)    ) or
	 (math.abs(pR) < osbf.cfg_learn_threshold) ) then
      osbf.learn(text, osbf.cfg_dbset, class_index, osbf.cfg_learn_flags)
      if (osbf.cfg_log_learned) then
        spamfilter_log(text, user_log_dir .. string.format("%s_%d.log",
							"learned_as_class",
							class_index)) 
      end
      return "", pR
    else
      return string.format(
	  "%s: score = %.2f; out of learning region: [-%.1f, %.1f]",
	  osbf.cfg_training_not_necessary,
	  pR, osbf.cfg_learn_threshold, osbf.cfg_learn_threshold)
    end
  else
    return msg_error, pR
  end
end

function process_command (cmd, cmd_args)
  local valid_command = true
  local orig_header, orig_msg
  local body = ""
  local training_result_subj = ""
  local sfid

  -- extract args, if any
  local arg1, arg2 = string.match(cmd_args, "^%s*(%S+)%s*(.*)")

  if (cmd == "learn" or cmd == "unlearn" or cmd == "recover" or
      cmd == "classify") then

    -- Remove content-type to avoid conflicts with the inserted body
    header = clean_header(header, {"Content%-Type:"})

    -- try to find the sfid to recover the original message
    if (arg2 and string.match(arg2, "sfid%-")) then
      sfid = arg2
    elseif (arg1 and string.match(arg1, "sfid%-")) then
      sfid = arg1
    end

    if (not sfid) then
      -- if the sfid was not given in the command line, extract it
      -- from the in-reply-to field in the reply
      sfid = string.match(header,
	 "\n[Ii][Nn]%-[Rr][Ee][Pp][Ll][Yy]%-[Tt][Oo]%s*:.-%((sfid%-.-)%)\n[^ \t]")

      -- if not found in the "In-Reply-To", try the "References" header
      if (not sfid) then
        sfid = string.match(header,
	 "\n[Re][Ee][Ff][Ee][Rr][Ee][Nn][Cc][Ee][Ss]%s*:.-%((sfid%-.-)%)\n[^ \t]")
      end
    end

    -- recover original header and body
    if sfid then
      local f = io.open(user_log_dir .. sfid, "r")
      if (f) then
        orig_msg = f:read("*all")
	orig_header = string.match(orig_msg, "^(.-\n\n)")
	body = " " .. orig_header
        f:close()
      else
        training_result_subj = "Original message not available any more " ..
				sfid
	body = [[
Copies of the original messages are kept for a few days on the
server for training purposes. Probably the message you are trying
to train with has already been deleted.

Try with a recent message, for instance one you received today,
to check the training mechanism. You may use the "recover" command,
instead of "learn" or "unlearn", for just checking if the message
is there.
]]
        return true, training_result_subj, body
      end
    else
        training_result_subj = "Spamfilter ID (SFID) not found"
	body = [[
The possible reasons why a SFID isn't found:
- You didn't use the "Reply" facility of your mail client;
- The message wasn't previously classified by the spamfilter;
- A bug in the spamfilter.

If you are sure about the first two possibilities, ask your
email system administrator for support.
]]
        return true, training_result_subj, body
    end

    if (cmd == "learn") then

       if (arg1 == "spam") then
         local msg, old_pR = osbf_train(orig_msg, osbf.cfg_spam_index)
	 pR, class_probs = osbf.classify(orig_msg, osbf.cfg_dbset,
                                         osbf.cfg_classify_flags)
	 if (msg == "") then

	   -- if we are still within learn_threshold and the previous
	   -- classification was a mistake or the pR improvement was
	   -- less than the unlearn_threshold, do an unlearning from
	   -- the opposite class
           if pR > -osbf.cfg_learn_threshold and
              (old_pR >= 0  or
	       (old_pR - pR) < osbf.cfg_unlearn_threshold) then
             osbf.unlearn(orig_msg, osbf.cfg_dbset, osbf.cfg_nonspam_index,
	     		    osbf.cfg_learn_flags)
	     pR, class_probs = osbf.classify(orig_msg, osbf.cfg_dbset,
                               osbf.cfg_classify_flags)
	   end

	   training_result_subj = osbf.cfg_trained_as_spam  ..
				  ": " .. string.format("%.2f", old_pR) ..
				  " -> " .. string.format("%.2f", pR)
	 else
	   training_result_subj = msg
	 end

       elseif (arg1 == "nonspam") then

         local msg, old_pR = osbf_train(orig_msg, osbf.cfg_nonspam_index)
	 pR, class_probs = osbf.classify(orig_msg, osbf.cfg_dbset,
                                         osbf.cfg_classify_flags)
	 if (msg == "") then

	   -- if we still are within learn_threshold and the previous
	   -- classification was a mistake or the pR improvement was
	   -- less than the unlearn_threshold, do an unlearning from
	   -- the opposite class
           if pR < osbf.cfg_learn_threshold and
              (old_pR < 0 or
	       (pR - old_pR) < osbf.cfg_unlearn_threshold) then
             osbf.unlearn(orig_msg, osbf.cfg_dbset, osbf.cfg_spam_index,
	     		osbf.cfg_learn_flags)
	     pR, class_probs = osbf.classify(orig_msg, osbf.cfg_dbset,
                               osbf.cfg_classify_flags)
	   end

	   training_result_subj = osbf.cfg_trained_as_nonspam  ..
				  ": " .. string.format("%.2f", old_pR) ..
				  " -> " .. string.format("%.2f", pR)
	 else
	   training_result_subj = msg
	 end

       else

         valid_command = false
       end

    elseif (cmd == "unlearn") then

       if (arg1 == "spam") then
	 old_pR, _ = osbf.classify(orig_msg, osbf.cfg_dbset,
                                         osbf.cfg_classify_flags)
         osbf.unlearn(orig_msg, osbf.cfg_dbset, osbf.cfg_spam_index,
	     		osbf.cfg_learn_flags)
	 pR, _ = osbf.classify(orig_msg, osbf.cfg_dbset,
                                         osbf.cfg_classify_flags)
         training_result_subj = "Message unlearned as spam: " .. 
				string.format("%.2f", old_pR) ..
				" -> " .. string.format("%.2f", pR)
				
       elseif (arg1 == "nonspam") then
	 old_pR, _ = osbf.classify(orig_msg, osbf.cfg_dbset,
                                         osbf.cfg_classify_flags)
         osbf.unlearn(orig_msg, osbf.cfg_dbset, osbf.cfg_nonspam_index,
			 osbf.cfg_learn_flags)
	 pR, _ = osbf.classify(orig_msg, osbf.cfg_dbset,
                                         osbf.cfg_classify_flags)
         training_result_subj = "Message unlearned as nonspam: " ..
				string.format("%.2f", old_pR) ..
				" -> " .. string.format("%.2f", pR)
       else
         valid_command = false
       end

    elseif (cmd == "recover") then

      boundary = string.gsub(sfid, "@.*", "=-=-=", 1)

      -- alter original header!
      header = clean_header(header, {"Content%-Type:"})
      header = string.gsub(header,
      "\n[Cc][Oo][Nn][Tt][Ee][Nn][Tt]%-[Tt][Yy][Pp][Ee]%s*:.-\n([^ \t])", "\n%1", 1)
      header = header .. "Content-Type: multipart/mixed;\n" ..
         " boundary=\"" .. boundary .. "\"\n"

      training_result_subj = "The recovered message is attached"


      -- protect and keep the original envelope-from line
      local xooef = ""
      if string.match(orig_msg, "^From ") then
        xooef = "X-OSBF-Original-Envelope-From:\n\t"
      end

      body = "--" .. boundary ..
         "\nContent-Type: message/rfc822;\n" ..
         " name=\"Recovered Message\"\n" ..
         "Content-Transfer-Encoding: 8bit\n" ..
         "Content-Disposition: inline;\n" ..
         " filename=\"Recovered Message\"\n\n" ..
         xooef .. orig_msg ..
         "\n--" .. boundary .. "--\n"

    elseif (cmd == "classify") then
      header = clean_header(header, {"Content%-Type:"})
      pR, class_probs = osbf.classify(orig_msg, osbf.cfg_dbset,
					 osbf.cfg_classify_flags)
      training_result_subj = "pR = " .. string.format("%.4f", pR)
    end
  elseif (cmd == "whitelist" or cmd == "blacklist") then
    header = clean_header(header, {"Content%-Type:"})
    if (arg1 == "add" or arg1 == "del") then
      local result, msg_err
      if arg1 == "add" then
        if cmd == "whitelist" then
          result, msg_err = list_insert(arg2, whitelist, cmd)
        else
          result, msg_err = list_insert(arg2, blacklist, cmd)
	end
      else
        if cmd == "whitelist" then
          result, msg_err = list_delete(arg2, whitelist, cmd)
        else
          result, msg_err = list_delete(arg2, blacklist, cmd)
	end
      end
      if not result then
        training_result_subj = "Error: " .. msg_err
      else
        training_result_subj = cmd .. " " .. arg1 .. " " .. arg2
        if cmd == "whitelist" then
          body = list_sprint(whitelist, "whitelist")
        else
          body = list_sprint(blacklist, "blacklist")
        end
      end
    elseif (arg1 == "show") then
      header = clean_header(header, {"Content%-Type:"})
      training_result_subj = cmd .. " " .. "contents"
      if cmd == "whitelist" then
        body = list_sprint(whitelist, "whitelist")
      else
        body = list_sprint(blacklist, "blacklist")
      end
    else 
      header = clean_header(header, {"Content%-Type:"})
      training_result_subj = "Unknown command"
      body =  cmd .. " <password> " .. cmd_args
    end
  elseif (cmd == "stats") then
    header = clean_header(header, {"Content%-Type:"})
    training_result_subj = "OSBF-Lua database statistics"
    body = ""
    for _, classfile in ipairs(osbf.cfg_dbset.classes) do
      if (body == "") then
        body = db_stats(classfile)
      else
        body = body .. "\n" .. db_stats(classfile)
      end
    end
  elseif (cmd == "help") then
    header = clean_header(header, {"Content%-Type:"})
    training_result_subj = "spamfilter help"
    body = help()
  else
    valid_command = false
  end

  return valid_command, training_result_subj, body
end

